package middleware

import (
	"context"
	utils "hft/Http/Utils"
	"hft/global"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
)

func NewRateLimiter(requestLimitMinute, requestLimitHour int, banDuration time.Duration) gin.HandlerFunc {
	return func(c *gin.Context) {
		clientIP := c.ClientIP()
		currentTime := time.Now()
		currentTimeUnix := currentTime.Unix()
		banEndTime := currentTime.Add(banDuration).Unix()
		banKey := "ban:" + clientIP

		// Check if the IP is permanently banned
		bannedKey := "banned:" + clientIP
		_, err := global.OC_REDIS.Get(bannedKey).Result()
		if err == nil {
			utils.FailWithMessage("该ip被永久封禁", c)
			c.Abort()
			return
		}

		// Check if the IP is temporarily banned
		banEndTimeStr, err := global.OC_REDIS.Get(banKey).Result()
		if err == nil {
			banEndTimeUnix, _ := strconv.ParseInt(banEndTimeStr, 10, 64)
			if banEndTimeUnix > currentTimeUnix {
				utils.FailWithMessage("ip被封禁到"+banEndTimeStr, c)
				c.Abort()
				return
			} else {
				// IP ban has expired, remove the ban key
				global.OC_REDIS.Del(banKey)
			}
		}

		// Count the number of requests for the IP in the past minute
		minuteKey := "requests:minute:" + clientIP
		minuteCount, err := global.OC_REDIS.Incr(minuteKey).Result()
		if err != nil {
			c.JSON(500, gin.H{"error": "服务器错误"})
			c.Abort()
			return
		}

		// Set an expiration time for the key
		if minuteCount == 1 {
			global.OC_REDIS.Expire(minuteKey, time.Minute)
		}

		// Check if the request limit for the minute is exceeded
		if minuteCount > int64(requestLimitMinute) {
			// Temporarily ban the IP
			global.OC_REDIS.Set(banKey, banEndTime, banDuration)
			utils.FailWithMessage("ip请求过于频繁", c)
			c.Abort()
			return
		}

		// Count the number of requests for the IP in the past hour
		hourKey := "requests:hour:" + clientIP
		hourCount, err := global.OC_REDIS.Incr(hourKey).Result()
		if err != nil {
			c.JSON(500, gin.H{"error": "服务器错误"})
			c.Abort()
			return
		}

		// Set an expiration time for the key
		if hourCount == 1 {
			global.OC_REDIS.Expire(hourKey, time.Hour)
		}

		// Check if the request limit for the hour is exceeded
		if hourCount > int64(requestLimitHour) {
			// Permanently ban the IP
			global.OC_REDIS.Set(bannedKey, "1", 0)
			utils.FailWithMessage("该ip已被封禁", c)
			c.Abort()

			// Store the ban record in MongoDB
			go func() {
				banRecord := map[string]interface{}{
					"ip":        clientIP,
					"ban_start": currentTime,
					"ban_end":   currentTime.Add(banDuration),
				}
				collection := global.OC_MON.Database("ip_db").Collection("ban_records")
				_, _ = collection.InsertOne(context.TODO(), banRecord)
			}()

			return
		}

		// Continue processing the request
		c.Next()
	}
}
